import os
import random


xmlfilepath = 'D:/资料/CV/video/VOCdevkit/VOC2007/Annotations'

total_xml = os.listdir(xmlfilepath)


from  sklearn.model_selection import train_test_split

trainval, test = train_test_split(total_xml, test_size=0.3, random_state=42)

train, val = train_test_split(trainval, test_size=0.142, random_state=42)

print(len(train),len(test),len(val))

# import  matplotlib.pyplot as plt
#
#
# # 每个场景的样本数
# plt.bar(['market','station','metro','street_night','street'],[len(market_xml),len(station_xml),len(metro_xml),len(street_night_xml),len(street_xml)])
# plt.title("Sample Scenario")# 图名
# plt.xlabel('Scene')
# plt.ylabel('Num Of Sample') #x轴
# plt.legend() # 图例
#
# for x, y in zip(['market','station','metro','street_night','street'], [len(market_xml),len(station_xml),len(metro_xml),len(street_night_xml),len(street_xml)]):
#     plt.text(x , y , '%d' % y, ha='center', va='bottom')
#
# plt.show()
#
#
#
# # 训练  测试  验证的划分情况
# plt.pie([len(train_xmls),len(val_xmls),len(test_xmls)],labels = ['train','val','test'],autopct='%1.1f%%',shadow=False,startangle=90)
# plt.title("Sample Scenario")# 图名
# plt.legend() # 图例
# plt.show()
#
# import matplotlib.pyplot as plt
#
#
# x = ['market','station','metro','street_night','street']
# train = [len(market_train),len(station_train),len(metro_train),len(street_night_train),len(street_train)]
# test = [len(market_test),len(station_test),len(metro_test),len(street_night_test),len(street_test)]
# val = [len(market_val),len(station_val),len(metro_val),len(street_night_val),len(street_val)]
#
# total_width, n = 0.8, 3
# # width = total_width / n
# # , width=width
# width =0.5
# plt.bar(x,  train, width=width, label='train')
# plt.bar(x , val, width=width,  label='val')
# plt.bar(x , test, width=width,  label='test')
# plt.xlabel('Scene')
# plt.ylabel('Num Of Sample') #x轴
# plt.legend()
#
# plt.show()



ftest = open('D:/资料/CV/video/VOCdevkit/VOC2007/ImageSets/Main/test.txt', 'w')
ftrain = open('D:/资料/CV/video/VOCdevkit/VOC2007/ImageSets/Main/train.txt', 'w')
fval = open('D:/资料/CV/video/VOCdevkit/VOC2007/ImageSets/Main/val.txt', 'w')


# #
# for train_xml in train:
#     name = train_xml[:-4] + '\n'
#     ftrain.write(name)
#
# for val_xml in val:
#     name = val_xml[:-4] + '\n'
#     fval.write(name)
#
#
# for test_xml in test:
#     name = test_xml[:-4] + '\n'
#     ftest.write(name)

ftrain.close()
fval.close()
ftest.close()